import dep_graph.keys
import dep_graph.cache_view
import dep_graph.pickled_database
import logging

logger = logging.getLogger(__name__)

class Error(Exception):
  """Root level exception"""

class KeyMissingException(Error):
  """Could not find key in graph."""
  def __init__(self, key):
    super().__init__()
    self.key = key

class FileBackedCachedGraph(dep_graph.pickled_database.PickledDatabase):
  def __init__(self, *args):
    super().__init__(*args)
    self.rules = {} 
    self.filesets = {}
    self.build_files = {}
    self.files = {}

  def create_gen_rule_for_key(self, key):
    """Creates a new dep_graph.cache_view.GenRule for key.

    Forcibly replaces any existing value for key in self.

    Args:
      key: dep_graph.keys.GenRuleKey The key for the new GenRule.

    Returns the newly created GenRule."""
    self.try_delete_gen_rule_with_key(key)
    try:
      rule = dep_graph.cache_view.GenRule(self, key)
      logger.debug("created rule: " + str(rule))
      self.write(key, rule)
      self.rules[key] = rule
      return rule
    except:
      self.try_delete_gen_rule_with_key(key)
      raise

  def get_gen_rule_with_key(self, key):
    """Retrieves the GenRule corresponding to key.

    Args:
      key: The dep_graph.keys.GenRuleKey to look up.

    Returns the dep_graph.cache_view.GenRule corresponding to key or None.
    """
    if key in self.rules:
      return self.rules[key]
    value = self.lookup(key)
    if value:
      assert isinstance(value, dep_graph.cache_view.GenRule)
      self.rules[key] = value
      return value

  def get_fset_with_key(self, key):
    """Retrieves the FileSet corresponding to key, creating it if necessary.

    Args:
      key: The dep_graph.keys.FileSetKey to look up.

    Returns the dep_graph.cache_view.FileSet corresponding to key.
    """
    if key in self.filesets:
      return self.filesets[key]
    fset = self.lookup(key)
    if fset:
      assert isinstance(fset, dep_graph.cache_view.FileSet)
      self.filesets[key] = fset
      return fset
    logger.debug("creating fset: " + str(key))
    fset = dep_graph.cache_view.FileSet(self, key)
    self.write(key, fset)
    self.filesets[key] = fset
    return fset

  def get_existing_fset_with_key(self, key):
    """Retrieves the FileSet corresponding to key or raises an error.

    Args:
      key: The dep_graph.keys.FileSetKey to look up.

    Returns the dep_graph.cache_view.FileSet corresponding to key.
    """
    if key in self.filesets:
      return self.filesets[key]
    fset = self.lookup(key)
    assert fset
    assert isinstance(fset, dep_graph.cache_view.FileSet)
    self.filesets[key] = fset
    return fset

  def get_fsets_for_file(self, fname):
    """Gets the FileSet corresponding to a file.

    Args:
      fname: The string name of the file to look up.

    Returns generator.<dep_graph.cache_view.FileSet> linked to fname."""
    fset_key_set = self.get_fset_keys_for_file(fname)
    for fset_key in tuple(fset_key_set):
      yield self.get_existing_fset_with_key(fset_key)

  def get_fset_keys_for_file(self, fname):
    """Gets the FileSet corresponding to a file.

    Args:
      fname: The string name of the file to look up.

    Returns set.<dep_graph.cache_view.FileSetKey> linked to fname."""
    key = dep_graph.keys.FileKey(fname)
    if key in self.files:
      return self.files[key]
    fset_key_set = self.lookup(key)
    if not fset_key_set:
      fset_key_set = set()
    self.files[key] = fset_key_set
    return fset_key_set
      
  def unlink_file_from_fset(self, fname, fset_key):
    """UnLinks a file (using a dep_graph.keys.FileKey) from a FileSet.

    Args:
      fname: The string name of the file to unlink.
      fset_key: The key of the dep_graph.cache_view.FileSet to unlink fname from.
    """
    file_key = dep_graph.keys.FileKey(fname)
    logger.debug("writing FileKey: " + str(file_key))
    
    fset_key_set = self.get_fset_keys_for_file(fname) 
    fset_key_set.remove(fset_key)

    if (len(fset_key_set)):
      self.write(file_key, fset_key_set)
    else:
      self.remove(file_key)

  def link_file_to_fset(self, fname, fset_key):
    """Links a file (using a dep_graph.keys.FileKey) to a FileSet.

    TODO(parker): Make this a set. Need to unlink old filesets.

    Args:
      fname: The string name of the file to link.
      fset_key: The key of the dep_graph.cache_view.FileSet to link fname to.
    """
    file_key = dep_graph.keys.FileKey(fname)
    logger.debug("writing FileKey: " + str(file_key))
    
    fset_key_set = self.get_fset_keys_for_file(fname) 
    try:
      fset_key_set.add(fset_key)
    finally:
      self.write(file_key, fset_key_set)

  def write_gen_rule(self, key):
    """Writes a dep_graph.cache_view.GenRule out to the file.

    Args:
      key: The key of the GenRule to write out.
    """
    rule = self.get_gen_rule_with_key(key)
    assert rule is not None
    self.write(key, rule)

  def write_fset(self, key):
    """Writes a dep_graph.cache_view.FileSet out to the file.

    Args:
      key: The key of the FileSet to write out.
    """
    fset = self.get_fset_with_key(key)
    assert fset is not None
    self.write(key, fset)

  def try_delete_gen_rule_with_key(self, key):
    """Removes the rule with the given key from self.

    It is not an error if there is no rule corresponding to key.

    Args:
      key: The key of the GenRule to remove.
    """
    if key in self.rules:
      rule = self.rules[key]
      del self.rules[key]
    else:
      rule = self.lookup(key)
    if rule:
      assert isinstance(rule, dep_graph.cache_view.GenRule)
      logger.debug("del rule: " + str(key))

      to_remove = set(rule.reverse_dep_keys())
      rule.remove_from_graph()
      for key in to_remove:
        self.try_delete_gen_rule_with_key(key)


  def delete_gen_rule_with_key(self, key):
    """Removes the rule with the given key from self or raises an exception.

    Args:
      key: The key of the GenRule to remove.
    """
    if key in self.rules:
      rule = self.rules[key]
      del self.rules[key]
    else:
      rule = self.lookup(key)
    if not rule:
      raise KeyMissingException(key)
    assert isinstance(rule, dep_graph.cache_view.GenRule)
    logger.debug("del rule: " + str(key))

    to_remove = set(rule.reverse_dep_keys())
    rule.remove_from_graph()
    for key in to_remove:
      self.try_delete_gen_rule_with_key(key)

  def delete_fset_with_key(self, key):
    """Removes the rule with the given key from self.

    Args:
      key: The key of the FileSet to remove.
    """
    if key in self.filesets:
      fset = self.filesets[key]
      del self.filesets[key]
    else:
      fset = self.lookup(key)
    assert fset
    assert isinstance(fset, dep_graph.cache_view.FileSet)
    logger.debug("del fset: " + str(key))
    fset.remove_from_graph()

  def delete_build_file(self, fname):
    """Removes the build file with the given filename from self.

    Args:
      fname: The filename of the build file to remove.

    Returns a sequence of the names of all the files that were removed.
    """
    key = dep_graph.keys.BuildFileKey(fname)
    if key in self.build_files:
      build = self.build_files[key]
      del self.build_files[key]
    else:
      build = self.lookup(key)
    if not build:
      return set()
    assert isinstance(build, set)
    self.remove(key)

    r = set()
    for rule_key in build:
      rule = self.get_gen_rule_with_key(rule_key)
      if not rule:
        continue
      for bucket in (rule.out, rule.meta):
        for fset_key in bucket.getFsetKeys():
          for fname in self.get_existing_fset_with_key(fset_key).files:
            r.add(fname)

    for rule_key in build:
      self.try_delete_gen_rule_with_key(rule_key)
    return r

  def add_gen_rule_to_build_file(self, fname, rule_key):
    """Adds a GenRule to a build file.

    Args:
      fname: The filename of the build file to add to.
      rule_key: The key of the GenRule to add.
    """
    key = dep_graph.keys.BuildFileKey(fname)
    if key in self.build_files:
      build = self.build_files[key]
    else:
      build = self.lookup(key)
      if not build:
        build = set()
      self.build_files[key] = build
    assert isinstance(build, set)
    build.add(rule_key)
    self.write(key, build)
    logger.debug("adding " + str(rule_key) + " to " + str(fname))
